# encoding: utf-8
"""
@author:  shaoqijun
@date:  2020/7/26
"""

import os
import cv2
from tqdm import tqdm
from config import cfg

class gen_img_list():
    def __init__(self,in_dir = r'D:\dataset\button_white', out_dir = r'D:\dataset\data\button'):
        self.in_dir = in_dir
        self.out_dir = out_dir
        self.img_list_txt = os.path.join(out_dir, "button_white_list.txt")
        self._check_before_run()
        self._write_txt(if_show=True)


    def _makedir(self,dir):
        if not os.path.exists(dir):
            os.makedirs(dir)

    def _check_before_run(self):
        self._makedir(self.out_dir)
        if not os.path.exists(self.in_dir):
            raise RuntimeError("'{}' is not available".format(self.in_dir))
        if not os.path.exists(self.out_dir):
            raise RuntimeError("'{}' is not available".format(self.out_dir))

    def _read_img(self,img_path):
        got_img = False
        while not got_img:
            try:
                img = cv2.imread(img_path)
                got_img = True
                cv2.imshow('img',img)
                cv2.waitKey(1)
            except IOError:
                print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(
                    img_path))
                pass

    def _write_txt(self, if_show=False):
        with open(self.img_list_txt,'w') as f:
            for image_name in tqdm(os.listdir(self.in_dir)):
                pwd = os.path.join(self.in_dir, image_name)
                f.writelines('{}\n'.format(pwd))  #image_name

                if if_show==True:
                    image_path = os.path.join(self.in_dir,image_name)
                    self._read_img(image_path)

if __name__ == '__main__':
    '''
    输入：     in_dir:图片数据目录
              out_dir:txt的输出路径
    输出：     img_list.txt:图片名字的列表txt，每行为图片名字
    '''
    in_dir = r'D:\dataset\button_white'
    out_dir = r'D:\dataset\data\button'
    #gen_img_list(in_dir=in_dir, out_dir=out_dir)
    gen_img_list(in_dir=cfg.IMGLIST.IMG_IN_DIR, out_dir=cfg.IMGLIST.IMG_OUT_DIR)


